/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_smooth.h"
#ifdef USE_OPENMP
#include <omp.h>
#endif

namespace YHAMG
{

void ParCSR_Smooth(const ParCSR_Matrix& A, const Vector& D_Recip, const Par_Vector& b, const Par_Vector& x, const Par_Vector& w, int type, double omega)
{
	int n = A.LocLoc.size[0];
	int Ext_size = A.LocExt.size[1];

	int* LocLoc_rowptr = A.LocLoc.rowptr;
	int* LocLoc_colind = A.LocLoc.colind;
	double* LocLoc_values = A.LocLoc.values;
	int* LocExt_rowptr = A.LocExt.rowptr;
	int* LocExt_colind = A.LocExt.colind;
	double* LocExt_values = A.LocExt.values;

	double* Dv_Recip = D_Recip.values;
	double* xv_Loc = x.Loc.values;
	double* bv_Loc = b.Loc.values;
	double* wv_Loc = w.Loc.values;
	
	double* xv_Ext = A.x_Ext.values;

	A.ExchangeHalo(x.Loc, A.x_Ext);

	if (type == 0)
	{
		for (int i = 0; i < n; ++i)
			wv_Loc[i] = xv_Loc[i];
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			double temp = bv_Loc[i];
			for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
				temp -= LocLoc_values[j] * wv_Loc[LocLoc_colind[j]];
			for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
				temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
			xv_Loc[i] += omega * Dv_Recip[i] * temp;
		}
	}
	else if (type == 1)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif
		if (num_threads == 1)
		{
			for (int i = 0; i < n; ++i)
			{
				double temp = bv_Loc[i];
				for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					temp -= LocLoc_values[j] * xv_Loc[LocLoc_colind[j]];
				for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
					temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
				xv_Loc[i] += omega * Dv_Recip[i] * temp;
			}
		}
		else
		{
			for (int i = 0; i < n; ++i)
				wv_Loc[i] = xv_Loc[i];
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;

				for (int i = begin; i < end; ++i)
				{
					double temp = bv_Loc[i];
					for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					{
						int jcol = LocLoc_colind[j];
						double jval = LocLoc_values[j];

						if (jcol < begin || jcol >= end)
							temp -= jval * wv_Loc[jcol];
						else
							temp -= jval * xv_Loc[jcol];
					}
					for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
						temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
					xv_Loc[i] += omega * Dv_Recip[i] * temp;
				}
			}
		}
	}
	else if (type == 2)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		if (num_threads == 1)
		{
			for (int i = n - 1; i >= 0; --i)
			{
				double temp = bv_Loc[i];
				for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					temp -= LocLoc_values[j] * xv_Loc[LocLoc_colind[j]];
				for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
					temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
				xv_Loc[i] += omega * Dv_Recip[i] * temp;
			}
		}
		else
		{
			for (int i = 0; i < n; ++i)
				wv_Loc[i] = xv_Loc[i];
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;

				for (int i = end - 1; i >= begin; --i)
				{
					double temp = bv_Loc[i];
					for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					{
						int jcol = LocLoc_colind[j];
						double jval = LocLoc_values[j];

						if (jcol < begin || jcol >= end)
							temp -= jval * wv_Loc[jcol];
						else
							temp -= jval * xv_Loc[jcol];
					}
					for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
						temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
					xv_Loc[i] += omega * Dv_Recip[i] * temp;
				}
			}
		}
	}
	else if (type == 3)
	{
		int num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		if (num_threads == 1)
		{
			for (int i = 0; i < n; ++i)
			{
				double temp = bv_Loc[i];
				for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					temp -= LocLoc_values[j] * xv_Loc[LocLoc_colind[j]];
				for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
					temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
				xv_Loc[i] += omega * Dv_Recip[i] * temp;
			}

			for (int i = n - 1; i >= 0; --i)
			{
				double temp = bv_Loc[i];
				for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					temp -= LocLoc_values[j] * xv_Loc[LocLoc_colind[j]];
				for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
					temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
				xv_Loc[i] += omega * Dv_Recip[i] * temp;
			}
		}
		else
		{
			for (int i = 0; i < n; ++i)
				wv_Loc[i] = xv_Loc[i];
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
			for (int t = 0; t < num_threads; ++t)
			{
				int begin = n * t / num_threads;
				int end = n * (t + 1) / num_threads;

				for (int i = begin; i < end; ++i)
				{
					double temp = bv_Loc[i];
					for (int j = LocLoc_rowptr[i]; j < LocLoc_rowptr[i + 1]; ++j)
					{
						int jcol = LocLoc_colind[j];
						double jval = LocLoc_values[j];

						if (jcol < begin || jcol >= end)
							temp -= jval * wv_Loc[jcol];
						else
							temp -= jval * xv_Loc[jcol];
					}
					for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
						temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
					xv_Loc[i] += omega * Dv_Recip[i] * temp;
				}

				for (int i = end - 1; i >= begin; --i)
				{
					double temp = bv_Loc[i];
					for (int j = LocLoc_rowptr[i + 1] - 1; j >= LocLoc_rowptr[i]; --j)
					{
						int jcol = LocLoc_colind[j];
						double jval = LocLoc_values[j];

						if (jcol < begin || jcol >= end)
							temp -= jval * wv_Loc[jcol];
						else
							temp -= jval * xv_Loc[jcol];
					}
					for (int j = LocExt_rowptr[i + 1] - 1; j >= LocExt_rowptr[i]; --j)
						temp -= LocExt_values[j] * xv_Ext[LocExt_colind[j]];
					xv_Loc[i] += omega * Dv_Recip[i] * temp;
				}
			}
		}
	}
}

}